# *_* : coding: utf-8 *_*

'''
datasets process for object detection project.
for convert customer dataset format to coco data format,
'''

import traceback
import argparse
import json
import cv2
import csv
import os

__CLASS__ = ['__background__', 'CatFace']   # class dictionary, background must be in first index.

def argparser():
    parser = argparse.ArgumentParser("define argument parser for pycococreator!")
    parser.add_argument("-r", "--image_root", default='D:\\catface\\train\\', help="path of root directory")
    parser.add_argument("-p", "--phase_folder", default=["train", "val"], help="datasets split")

    return parser.parse_args()

def MainProcessing(args):
    '''main process source code.'''
    annotations = {}    # annotations dictionary, which will dump to json format file.
    image_root = args.image_root
    phase_folder = args.phase_folder
    with open('train.csv', 'r') as f:
        reader = csv.reader(f)
        result = list(reader)
        result = result[1:]

    # coco annotations info.
    annotations["info"] = {
        "description": "customer dataset format convert to COCO format",
        "url": "http://cocodataset.org",
        "version": "1.0",
        "year": 2020,
        "contributor": "ezra",
        "date_created": "2020/03/15"
    }
    # coco annotations licenses.
    annotations["licenses"] = [{
        "url": "https://www.apache.org/licenses/LICENSE-2.0.html",
        "id": 1,
        "name": "Apache License 2.0"
    }]
    # coco annotations categories.
    annotations["categories"] = []
    for cls, clsname in enumerate(__CLASS__):
        if clsname == '__background__':
            continue
        annotations["categories"].append(
            {
                "id": cls,
                "name": clsname,
                "supercategory": "Cat",
            }
        )
        for catdict in annotations["categories"]:
            if "CatFace" == catdict["name"]:
                catdict["keypoints"] = [0, 1, 2, 3, 4, 5, 6, 7, 8]
                catdict["skeleton"] = [[0,1],[1,2],[0,2],[3,4],[4,5],[5,6],[6,7],[7,8],[8,3]]

    for phase in phase_folder:
        annotations["images"] = []
        annotations["annotations"] = []
        fphase = open(phase + '.txt', 'r')
        step = 0
        for id, line in enumerate(fphase.readlines()):
            line = line.strip("\n")
            file_name = line + '.jpg'
            images_id = int(line)
            height, width, _ = cv2.imread(image_root + file_name).shape
            v = [2, 2, 2, 2, 2, 2, 2, 2, 2]
            point_str = result[images_id][1:]
            point = [int(k) for k in point_str]
            for j in range(9):
                if min(point[2*j:2*j+2]) < 0:
                    v[j] = 1
            keypoint = [point[0],  point[1],  v[0], point[2],  point[3],  v[1],
                        point[4],  point[5],  v[2], point[6],  point[7],  v[3],
                        point[8],  point[9],  v[4], point[10], point[11], v[5],
                        point[12], point[13], v[6], point[14], point[15], v[7],
                        point[16], point[17], v[8]]
            bw = max(point[0::2]) - min(point[0::2]) + 10
            bh = max(point[1::2]) - min(point[1::2]) + 10
            if (min(point[0::2]) - 5) < 0:
                x1 = 0
            else:
                x1 = (min(point[0::2]) - 5)
            if (min(point[1::2]) - 5) < 0:
                y1 = 0
            else:
                y1 = (min(point[1::2]) - 5)
            annotations["images"].append(
                {
                    "file_name": file_name,
                    "height": height,
                    "width": width,
                    "id": images_id
                }
            )
            # coco annotations annotations.
            annotations["annotations"].append(
                {
                    "id": id + 1,
                    "num_keypoints": 9,
                    "keypoints": keypoint,
                    "area": bw * bh,
                    "iscrowd": 0,
                    "image_id": images_id,
                    "bbox": [x1, y1, bw, bh],
                    "category_id": 1,
                    "segmentation": [],
                }
            )
            step += 1
            if step % 100 == 0:
                print("processing {} ...".format(step))

        json_path = phase+".json"
        with open(json_path, "w") as f:
            json.dump(annotations, f)


if __name__ == "__main__":
    print("begining to convert customer format to coco format!")
    args = argparser()
    try:
        MainProcessing(args)
    except Exception as e:
        traceback.print_exc()
    print("successful to convert customer format to coco format")
